import json

import cv2
import numpy as np
import scipy
import scipy.misc
import scipy.ndimage
import scipy.signal
import logging
import os.path
from logging import debug, info, error, warning
import math
import random
import MA.candate as ma
import pbcvt
import caffe

TMP_OUT = 'static/diag_result/'
if not os.path.exists(TMP_OUT):
    os.mkdir(TMP_OUT)
FORMAT = '[%(levelname)-5s]%(asctime)-8s %(filename)s:%(lineno)d %(message)s'
DATEFORMAT = '%H:%M:%S'
logging.basicConfig(level=logging.DEBUG, format=FORMAT, datefmt=DATEFORMAT)
logging.debug('start')
from celery import Celery

app = Celery('tasks', broker='redis://localhost:6379/0', backend='redis://localhost')

# load the model
net = caffe.Net('MA/bvlc_alexnet/deploy.prototxt',
                'MA/alexnet/caffe_alexnet_train_iter_27745.caffemodel',
                caffe.TEST)

blob = caffe.proto.caffe_pb2.BlobProto()
data = open('MA/data_mean.binaryproto', 'rb').read()
blob.ParseFromString(data)
arr = np.array(caffe.io.blobproto_to_array(blob))

# load input and configure preprocessing
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
transformer.set_mean('data', arr[0].mean(1).mean(1))
transformer.set_transpose('data', (2, 0, 1))
transformer.set_channel_swap('data', (2, 1, 0))
transformer.set_raw_scale('data', 255.0)


def imread(imgname):
    """Read image file form predifined file, To load different
    database, just need to change this function

    Args:
        imagename: name of the file, not the whole path
    Returns:
        Image read from the given name
    """
    img = scipy.misc.imread(imgname)
    # size = np.floor(np.array(img.shape) * 0.7).astype(np.int32)
    # img = scipy.misc.imresize(img, size)
    img = img / 255
    return img


def imsave(image_name, img, scale=False):
    """Save image file to folder defined in TMP_OUT
    Args:
        image_name: the output file name
        img: byte array of image
    Return:
        The path of the image
    """
    assert type(image_name) == str
    logging.debug('save to ' + image_name)
    if not scale:
        img[img < 0] = 0
        img[img > 1] = 1
    image_path = TMP_OUT + image_name + '.png'
    scipy.misc.imsave(image_path, img)
    return image_path


def imsave_with_index(image_name, img, original_name='', scale=True):
    tail = os.path.split(original_name)[1]
    name = os.path.splitext(tail)[0]
    file_name = '%s_%s' % (name, image_name)
    return imsave(file_name, img, scale)


def get_out_file_name(image_name, original_name):
    tail = os.path.split(original_name)[1]
    name = os.path.splitext(tail)[0]
    return '%s/%s_%s.png' % (TMP_OUT, name, image_name)


# def vessel_seg(image):
#     return pbcvt.vesselseg(image) > 0
#     # return image

def zoom(img: np.array, shape: np.array) -> np.array:
    if not isinstance(shape, (list, tuple, np.ndarray)):
        shape = [
            math.floor(img.shape[0] / img.shape[1] * shape),
            shape
        ]
        if len(img.shape) == 3:
            shape.append(img.shape[2])
    zoom_factor = [
        shape[0] / img.shape[0],
        shape[1] / img.shape[1],
    ]
    if len(img.shape) == 3:
        zoom_factor.append(1)
    return scipy.ndimage.zoom(img, zoom_factor, mode='nearest')


def line_dector(image):
    assert len(image.shape) == 2
    orishape = image.shape

    # image = zoom(image, 800)
    def get_line_response(image, detector_length=41):
        """
        detect lines in image
        :param detector_length: length of line detector
        :param image: original gray image
        :return: line map
        """
        center = math.floor(detector_length / 2)
        line_masks = []
        angle_list = [x * math.pi / 12 for x in range(12)]
        # angle_list = [math.pi / 2]
        for theta in angle_list:
            mask = np.zeros((detector_length, detector_length))
            dcol_unit = math.cos(theta)
            drow_unit = math.sin(theta)
            dcol = 0
            drow = 0
            while True:
                if math.fabs(dcol) >= center or math.fabs(drow) >= center:
                    break
                mask[center + math.floor(drow), center + math.floor(dcol)] = 1
                mask[center - math.floor(drow), center - math.floor(dcol)] = 1
                dcol += dcol_unit
                drow += drow_unit
            # mask = mask / mask.sum()
            line_masks.append(mask)
        line_response = np.zeros(image.shape)
        for mask in line_masks:
            # np.maximum(line_response, scipy.signal.convolve2d(image, mask, mode='same'), line_response)
            # line_response += scipy.signal.convolve2d(image, mask, mode='same')
            line_response += scipy.ndimage.convolve(image * 255, mask) > 120
        # line_response[line_response > 6*255] = 0
        return np.abs(line_response)

        # lines = np.zeros(image.shape)
        # for line_length in range(21,21,2):
        # lines = np.maximum(lines, get_line_response(image,line_length))
        # lines += get_line_response(image,line_length)

    return get_line_response(image, 91)


def get_mask(img):
    m = 1 - (img < 0.06)
    return scipy.ndimage.binary_erosion(m, np.ones((11, 11)), 10)


def get_ma(file_name):
    size = 85
    half_width = math.floor(size / 2)
    img = scipy.misc.imread(file_name)
    img = zoom(img,1500)
    pix_list = ma.get_local_max(img / 255)
    debug(len(pix_list))
    plist = []
    for index in range(len(pix_list)):
        crow = int(pix_list[index, 0])
        ccol = int(pix_list[index, 1])
        if not ((0 <= crow - half_width < crow + half_width < img.shape[0]) and
                    (0 <= ccol - half_width < ccol + half_width < img.shape[0])):
            # warning('index ' + str((row, col)) + ' out of range')
            continue

        # debug(img.shape)
        # debug(crow)
        # debug(ccol)
        # debug(type(crow))
        # debug((crow,ccol))
        patch = img[(crow - half_width):(crow + half_width), (ccol - half_width):(ccol + half_width), :]
        # debug(patch.shape)
        net.blobs['data'].data[...] = transformer.preprocess('data', patch)
        out = net.forward()['prob'][0].argmax()
        if out == 1:
            plist.append((crow, ccol))
    new_img = ma.drawPoint(img, plist)
    return plist, new_img


def gen_diagnose(fname, name='', age=-1, number='', gender=''):
    img = imread(fname)
    img = zoom(img, 1200)
    original_file_name = imsave_with_index('100original', img, fname)
    inv_green = 1 - img[:, :, 1]
    red = img[:, :, 0]

    mask = get_mask(red)
    # imsave_with_index('900_mask',mask,fname)

    # imsave_with_index('105red', red, fname)
    # imsave_with_index('110inv_green', inv_green, fname)
    # get disks of fundus
    disk = red > red.max() * 0.80
    maskwidth = 80
    disk[math.floor(disk.shape[0] / 2) - maskwidth:math.floor(disk.shape[0] / 2) + maskwidth,
    math.floor(disk.shape[1] / 2) - maskwidth:math.floor(disk.shape[1] / 2) + maskwidth] = True
    # imsave_with_index('100disk', disk, fname)

    vessel_seg = (pbcvt.vesselseg(img, 0) > 0) + 0
    vessel_seg = scipy.ndimage.binary_dilation(vessel_seg, np.ones((3, 3)), iterations=3)
    # imsave_with_index('101_vessel', vessel_seg, fname)

    # calculate background
    small_img_for_bg = zoom(inv_green, 600)
    smallbg_mean_filter = scipy.ndimage.filters.median_filter(small_img_for_bg, (41, 41), mode='nearest')
    bg = zoom(smallbg_mean_filter, inv_green.shape)
    assert bg.shape[0] == inv_green.shape[0]
    assert bg.shape[1] == inv_green.shape[1]
    # imsave_with_index('120bg', bg, fname)

    # calculate normed pic
    normed = inv_green - bg
    maxval = normed.max()
    minval = normed.min()
    # imsave_with_index('105normed', normed, fname)

    ha_candidate = (normed > (maxval * 0.14)) + 0
    ha_candidate *= (1 - vessel_seg)
    ha_candidate = scipy.ndimage.binary_closing(ha_candidate, np.ones((5, 5))) + 0
    ha_candidate *= (1 - disk)
    ha_candidate *= mask
    ha_candidate = scipy.ndimage.median_filter(ha_candidate, (5, 5))
    ex_candidate = (normed < minval * 0.13) + 0
    ex_candidate *= (1 - disk)
    ex_candidate *= mask
    # ha_img_name_raw = imsave_with_index('105ha_raw', ha_candidate, fname)
    # ex_img_name_raw = imsave_with_index('202ex_raw', ex_candidate, fname)
    area = mask.sum()
    retval, labels, stats, centroids = cv2.connectedComponentsWithStats(ha_candidate.astype(np.uint8))
    debug(cv2.CC_STAT_AREA)
    debug(retval)
    debug(stats)
    debug(centroids)

    # -----------------------------------------MA------------------------------------
    ma_list, ma_img = get_ma(fname)
    ma_counter = len(ma_list)
    debug('MA,FINISHED ............................................................')
    # -------------------------------------------------------------------------------
    # import pickle
    # import io
    # f = io.FileIO('model.mod', 'rb')
    # svc = pickle.load(f)
    # level = int(svc.predict([[ex_candidate.sum(), ha_candidate.sum(), ma_counter]])[0])
    level = 0
    if ma_counter > 5:
        level = max(level, 1)

    block_lesion_thresh = 25
    retval, labels, stats, centroids = cv2.connectedComponentsWithStats(ha_candidate.astype(np.uint8))
    has_ha = False
    for i in range(1, retval):
        if stats[i, cv2.CC_STAT_AREA] > block_lesion_thresh:
            level = max(level, 2)
            has_ha = True

    retval, labels, stats, centroids = cv2.connectedComponentsWithStats(ex_candidate.astype(np.uint8))
    for i in range(1, retval):
        if stats[i, cv2.CC_STAT_AREA] > block_lesion_thresh:
            level = max(level, 2)
        if stats[i, cv2.CC_STAT_AREA] > block_lesion_thresh and has_ha:
            level = max(level, 3)

    # -------------------------------------------------------------------------------
    # get images that illustrate the result

    ex_img_show = img.copy()
    ha_img_show = img.copy()

    alpha = 0.2
    ex_img_show[ex_candidate > 0, 1] *= alpha
    ex_img_show[ex_candidate > 0, 1] += 1 - alpha
    ha_img_show[ha_candidate > 0, 1] *= alpha
    ha_img_show[ha_candidate > 0, 1] += 1 - alpha
    ex_img_name = imsave_with_index('ex_img', ex_img_show, fname)
    ha_img_name = imsave_with_index('ha_img', ha_img_show, fname)
    ma_img_name = imsave_with_index('ma_img', ma_img, fname)

    # get the text result
    def cal_number(img, minthresh):
        counter = 0
        retval, label, stats, centroids = cv2.connectedComponentsWithStats(img.astype(np.uint8))
        for i in range(1, retval):
            if stats[i][cv2.CC_STAT_AREA] > minthresh:
                counter += 1
        return counter

    def get_location(img):
        assert len(img.shape) == 2
        r, c = 0, 0
        counter = 0
        for row in range(img.shape[0]):
            for col in range(img.shape[1]):
                if img[row, col]:
                    r += row
                    c += col
                    counter += 1
        if counter == 0:
            return ''
        r /= counter
        c /= counter
        r /= img.shape[0]
        c /= img.shape[1]
        retval = ''
        if c < 0.375:
            retval += "视乳头"
        elif c < 0.625:
            retval += "黄斑区"
        else:
            retval += "后极部"
            return retval
        if r > 0.5:
            retval += "下方"
        else:
            retval += "上方"
        return retval

    def cal_strenth(cimg, thresh1=20, thresh2=50):
        assert len(cimg.shape) == 2
        img = cimg
        retval, label, stats, centroids = cv2.connectedComponentsWithStats(img.astype(np.uint8))
        area = 0
        for i in range(1, retval):
            area += stats[i][cv2.CC_STAT_AREA]
        area /= retval
        if area < 5:
            return "未见"
        elif area < thresh1:
            return "疑似"

        elif area < thresh2:
            return "散在"
        else:
            return "出现"

    desp = ''
    desp += get_location(ha_candidate) + cal_strenth(ha_candidate, 40, 80) + "出血" + ','
    desp += get_location(ex_candidate) + cal_strenth(ex_candidate, 10, 80) + "渗出" + '。'

    result = {
        "ex_area": float(ex_candidate.sum()),
        "ex_ratio": float(ex_candidate.sum() / area),
        # "ex_img_raw": ex_img_name_raw,
        "ex_img": ex_img_name,
        "ha_area": float(ha_candidate.sum()),
        "ha_ratio": float(ha_candidate.sum() / area),
        # 'ha_img_raw': ha_img_name_raw,
        'ha_img': ha_img_name,
        'original_img': original_file_name,
        "ma_count": ma_counter,
        "ma_img": ma_img_name,
        'degree': float(level),
        'text': desp
    }
    print(result)

    # make pdf
    if type(result) is str:
        result = json.loads(result)
    from jinja2 import Environment, PackageLoader
    import jinja2
    import time
    env = Environment(
        loader=jinja2.FileSystemLoader(searchpath='./')
        # autoescape=jinja2.select_autoescape(['html', 'xml'])
    )
    dataTable = ""
    tableTemplate = "{} & {} & {} & {} & {} & {} \\\\ \\hline \n"
    size = ha_candidate.shape[0] * ha_candidate.shape[1]
    dataTable += tableTemplate.format(
        "出血",
        "-",
        area,
        (area / size * 100),
        get_location(ha_candidate),
        cal_strenth(ha_candidate, 40, 80)
    )
    dataTable += tableTemplate.format(
        "渗出",
        cal_number(ex_candidate, 10),
        area,
        (area / size * 100),
        get_location(ex_candidate),
        cal_strenth(ex_candidate, 10, 80)
    )

    finalresult = ""
    if level == 0:
        finalresult = "未见明显DR"
    elif level == 1:
        finalresult = "轻度NPDR"
    elif level == 2:
        finalresult = "中度NPDR"
    elif level == 3:
        finalresult = "重度NPDR"
    else:
        error('level not known')

    # temp = env.get_template('template.tex')
    # env.globals["name"] = name
    # env.globals["age"] = age
    # env.globals["time"] = time.ctime()
    # env.globals["number"] = number
    # env.globals["gender"] = gender
    # env.globals["img"] = "{" + original_file_name + "}"
    # env.globals["datatable"] = dataTable
    # env.globals["seg1"] = desp
    # env.globals["seg2"] = finalresult
    # env.globals["seg3"] = "定期随访"
    # report = temp.render()
    # tempFileName = '/tmp/' + '%09d.tex' % random.randint(0, 100000)
    # open(tempFileName, 'w').write(report)

    return json.dumps(result)


def run_worker(r):
    """
    this function runs the worker threads
    :param r: redis connection
    :return:
    """
    while True:
        fname = r.blpop('diagnose_tasks')
        result = gen_diagnose(fname)
        r.rpush('finished queue', result)


def run_test():
    dbroot = '/home/d/data/DME/'
    files = os.listdir(dbroot)
    for file_name in files:
        full_path = dbroot + file_name
        debug(full_path)
        gen_diagnose(full_path)


if __name__ == '__main__':
    run_test()
    # gen_diagnose('/home/d/disk/2016年1-10/1/C0011506.jpg')
    # import redis
    #
    # pool = redis.ConnectionPool(host='localhost', port=6379, db=0)
    # r = redis.Redis(connection_pool=pool)
    # run_worker(r)
